[BugFix][V0] Fix AssertionError for prompt_logprobs #19844

xu-song · 2025-06-19T07:25:02Z

Essential Elements of an Effective PR Description Checklist

The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
The test plan, such as providing test command.
The test results, such as pasting the results comparison before and after, or e2e results
(Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

Purpose

  File "/workspace/home/lab/.conda/envs/grpo/lib/python3.12/site-packages/vllm/model_executor/layers/sampler.py", line 956, in get_logprobs
    assert len(next_token_ids) == len(query_indices)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError

Reproduce of `prompt_logprobs`

import os
from vllm import LLM, SamplingParams
from transformers import AutoTokenizer

os.environ["VLLM_USE_V1"] = "0"

model_path = "/workspace/czy/model_weights/Meta-Llama-3.1-8B-Instruct"
model = LLM(model_path, tensor_parallel_size=1, enable_prefix_caching=True)

# set prompt_logprobs
sampling_params = SamplingParams(prompt_logprobs=1, max_tokens=10)
tokenizer = AutoTokenizer.from_pretrained(model_path)

chat_prompts = tokenizer.apply_chat_template([[{"role": "system", "content": "You are a helpful assistant"}, {"role": "user", "content": "Test 1"}]], tokenize=False)
output = model.generate(chat_prompts, sampling_params, use_tqdm=False)
print("OK")

# AssertionError in prefix cache
chat_prompts = tokenizer.apply_chat_template([[{"role": "system", "content": "You are a helpful assistant"}, {"role": "user", "content": "Test 2"}]], tokenize=False)
output = model.generate(chat_prompts, sampling_params, use_tqdm=False) # ERROR!

github-actions · 2025-06-19T07:25:13Z

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

gemini-code-assist

Summary of Changes

Hello @xu-song, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request addresses a critical AssertionError that occurred during log probability calculation within the sampler.py module. The fix involves a precise adjustment to an indexing calculation, resolving an off-by-one issue that was likely causing mismatches in token list lengths, thereby preventing the assertion failure.

Highlights

Bug Fix: Indexing Correction: Corrected an off-by-one error in the calculation of next_token_index_end within the _get_next_prompt_tokens function in vllm/model_executor/layers/sampler.py. This change removes an extraneous +1 from the end index calculation on line 1200.
Assertion Error Resolution: This fix directly addresses and resolves the AssertionError (Fixes #13105) that occurred when len(next_token_ids) did not match len(query_indices) during log probability calculation, by ensuring correct token indexing.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request fixes an AssertionError in the V0 sampler's prompt log probability calculation. The change correctly adjusts the slicing logic for next_prompt_tokens to ensure its length matches the number of query_indices, resolving an off-by-one error. The reasoning behind the fix seems sound based on how prompt logprobs are handled.

To further improve the robustness of this area, consider adding a specific unit test that reproduces the failing scenario (e.g., a prompt of a certain length processed with prompt_logprobs enabled, potentially with chunked prefill if that was a factor) to ensure this bug does not regress in the future. The PR description checklist also indicates that the test plan and results are not yet provided, which would be valuable for verifying the fix.

gemini-code-assist · 2025-06-19T07:27:50Z

vllm/model_executor/layers/sampler.py

+    next_token_index_end = min(computed_len + query_len,
                               len(prompt_tokens))


This change correctly adjusts the next_token_index_end to fix an off-by-one error when calculating the next_prompt_tokens for prompt log probabilities in the V0 sampler. Here's a breakdown of why this is the correct fix:

Context: This function _get_next_prompt_tokens is called by _get_prompt_logprobs (in the V0 sampler) to get the target tokens for which log probabilities are being computed.

query_len Argument: The query_len argument passed to this function from _get_prompt_logprobs corresponds to seq_group.query_len. This represents the total number of tokens in the current part of the prompt being processed (e.g., the full prompt length or a chunk in chunked prefill).

query_indices Length: The query_indices (which are seq_group.prompt_logprob_indices in _get_prompt_logprobs) represent the specific positions in the logits tensor for which prompt logprobs are needed. If the current prompt part has N tokens (P_0, P_1, ..., P_{N-1}), logprobs are typically computed for predicting P_1 (from context P_0), P_2 (from context P_0,P_1), and so on, up to P_{N-1}. Thus, there are N-1 such logprobs, and len(query_indices) is N-1 (this corresponds to prompt_logprob_len = seq_group.query_len - 1, assuming do_sample=True and num_prefill_sample=1 in _prepare_seq_groups).

next_token_ids Length: The next_token_ids (which are next_prompt_tokens returned by this function) must match len(query_indices). These are the target tokens P_1, ..., P_{N-1}.

Original Bug:

next_token_index_start = computed_len + 1.

Original next_token_index_end was effectively computed_len + query_len + 1 (before min).

This resulted in len(next_prompt_tokens) = (computed_len + query_len + 1) - (computed_len + 1) = query_len (which is N in the example above).

The assertion len(next_token_ids) == len(query_indices) became N == N-1, which failed, causing the AssertionError.

Fix:

New next_token_index_end is effectively computed_len + query_len (before min).

This results in len(next_prompt_tokens) = (computed_len + query_len) - (computed_len + 1) = query_len - 1 (which is N-1 in the example above).

The assertion now becomes N-1 == N-1, which passes.

This change ensures that the number of target tokens extracted matches the number of logit entries for which prompt logprobs are being calculated, resolving the AssertionError.

Signed-off-by: Xu Song <xusong.vip@gmail.com>

njhill · 2025-06-19T21:31:24Z

@xu-song what version of vLLM are you using in the provided example? This change only applies to V0, so I don't think it should apply to your example unless you are explicitly disabling V1 via the env var.

xu-song · 2025-06-20T01:34:39Z

latest V0

xu-song requested review from zhuohan123, youkaichao, alexm-redhat, comaniac and njhill as code owners June 19, 2025 07:25

gemini-code-assist bot reviewed Jun 19, 2025

View reviewed changes

xu-song changed the title ~~Update sampler.py~~ Fix AssertionError for prompt_logprobs Jun 19, 2025

gemini-code-assist bot reviewed Jun 19, 2025

View reviewed changes

xu-song changed the title ~~Fix AssertionError for prompt_logprobs~~ [BugFix] Fix AssertionError for prompt_logprobs Jun 19, 2025

xu-song force-pushed the patch-7 branch 2 times, most recently from 55f5364 to d34c43b Compare June 19, 2025 08:11

update

433f3d8

Signed-off-by: Xu Song <xusong.vip@gmail.com>

xu-song force-pushed the patch-7 branch from d34c43b to 433f3d8 Compare June 19, 2025 08:13

njhill added the v0 label Jun 19, 2025

xu-song changed the title ~~[BugFix] Fix AssertionError for prompt_logprobs~~ [BugFix][V0] Fix AssertionError for prompt_logprobs Jun 20, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Uh oh!

[BugFix][V0] Fix AssertionError for prompt_logprobs #19844

[BugFix][V0] Fix AssertionError for prompt_logprobs #19844

Uh oh!

xu-song commented Jun 19, 2025 •

edited by github-actions bot

Loading

Uh oh!

github-actions bot commented Jun 19, 2025

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot Jun 19, 2025

Uh oh!

njhill commented Jun 19, 2025

Uh oh!

xu-song commented Jun 20, 2025

Uh oh!

Uh oh!

		next_token_index_end = min(computed_len + query_len,
		len(prompt_tokens))

Uh oh!

[BugFix][V0] Fix AssertionError for prompt_logprobs #19844

Are you sure you want to change the base?

[BugFix][V0] Fix AssertionError for prompt_logprobs #19844

Uh oh!

Conversation

xu-song commented Jun 19, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Essential Elements of an Effective PR Description Checklist

Purpose

Reproduce of prompt_logprobs

Uh oh!

github-actions bot commented Jun 19, 2025

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Jun 19, 2025

Choose a reason for hiding this comment

Uh oh!

njhill commented Jun 19, 2025

Uh oh!

xu-song commented Jun 20, 2025

Uh oh!

Uh oh!

xu-song commented Jun 19, 2025 •

edited by github-actions bot

Loading

Reproduce of `prompt_logprobs`